#ifndef PW_PBL_BITFIELD_HPP
#define PW_PBL_BITFIELD_HPP

#include <type_traits>
#include <utility>
#include "RangedInt.hpp"
#include "Math.hpp"
#include "Numbers.hpp"
#include "RustMimic.hpp"
#include "BitEnum.hpp"
#include "ForceAlign.hpp"

namespace pw_pbl::inline bitfield
{
	template <typename... UNDER> requires(sizeof...(UNDER) <= 1)
	struct BitFieldStruct;
}

namespace pw_pbl::_bitfield
{
	template <typename T>
	struct FldTraits;

	template <std::integral T>
	struct FldTraits<T>
	{
		using raw = T;
		static constexpr auto EBits = Math::BitCntOf<T>;
		static constexpr T raw_to_fld(raw v) { return v; }
		static constexpr raw fld_to_raw(T v) { return v; }
	};

	template <>
	struct FldTraits<bool>
	{
		using raw = bool;
		static constexpr auto EBits = 1;
		static constexpr bool raw_to_fld(raw v) { return v; }
		static constexpr raw fld_to_raw(bool v) { return v; }
	};

	template <>
	struct FldTraits<NBool>
	{
		using raw = bool;
		static constexpr auto EBits = 1;
		static constexpr NBool raw_to_fld(raw v) { return NBool::new_with_raw(v); }
		static constexpr raw fld_to_raw(NBool v) { return v.raw; }
	};

	template <ccDuckEnum T>
	struct FldTraits<T>
	{
		using raw = pw_pbl::underlying_type_t<T>;
		static constexpr auto EBits = Math::BitCntOf<raw>;
		static constexpr T raw_to_fld(raw v) { return T{v}; }
		static constexpr raw fld_to_raw(T v) { return pw_pbl::to_underlying(v); }
	};

	template <ccRangedInt T>
	struct FldTraits<T>
	{
		using raw = typename T::store_type;
		static constexpr auto EBits = T::EBits;
		static constexpr T raw_to_fld(raw v) { return T{.store = v}; }
		static constexpr raw fld_to_raw(T v) { return v.store; }
	};

	template <std::derived_from<BitFieldStruct<>> T>
	struct FldTraits<T>
	{
		using raw = T::underlying_type;
		static constexpr auto EBits = Math::BitCntOf<raw>;
		static constexpr T raw_to_fld(raw v) { return T{v}; }
		static constexpr raw fld_to_raw(T v) { return v.raw; }
	};

	template <typename FLD>
	concept ccAsBitField = requires { typename FldTraits<FLD>::raw; };

	/// 位域引用中底层类型的特征类
	template <typename T>
	struct UnderTraits;

	template <std::unsigned_integral UNDER>
	struct UnderTraits<UnalignRef<UNDER>>
	{
		using raw_ref_type = UnalignRef<UNDER>;
		using struct_type = void;
		using underlying_type = std::decay_t<UNDER>;
		static constexpr bool EIsConst = std::is_const_v<std::remove_reference_t<UNDER>>;
		static constexpr underlying_type read(raw_ref_type ref) { return ref.read(); }
		static constexpr void write(raw_ref_type ref, underlying_type v) { ref.write(v); }
		static constexpr void alter(raw_ref_type ref, std::invocable<underlying_type &> auto func) { ref.alter(std::move(func)); }
		static constexpr auto *raw_ptr(raw_ref_type ref) { return ref.ptr; }
	};

	template <std::unsigned_integral UNDER>
	struct UnderTraits<UNDER> : UnderTraits<UnalignRef<UNDER &>>
	{
	};

	template <typename STRUCT> requires std::derived_from<std::decay_t<STRUCT>, BitFieldStruct<>>
	struct UnderTraits<STRUCT>
	{
		using raw_ref_type = std::conditional_t<std::is_rvalue_reference_v<STRUCT>, const std::remove_reference_t<STRUCT> &, STRUCT &>;
		using struct_type = std::decay_t<STRUCT>;
		using underlying_type = typename struct_type::underlying_type;
		static constexpr bool EIsConst = std::is_const_v<std::remove_reference_t<raw_ref_type>>;

		static constexpr underlying_type read(raw_ref_type ref) { return MaybeUnalignTraits<>::read(ref.raw); }
		static constexpr void write(raw_ref_type ref, underlying_type v) { MaybeUnalignTraits<>::write(ref.raw, v); }
		static constexpr void alter(raw_ref_type ref, std::invocable<underlying_type &> auto func) { MaybeUnalignTraits<>::alter(ref.raw, std::move(func)); }
		static constexpr auto *raw_ptr(raw_ref_type ref) { return &ref; }
	};

	template <size_t... START>
	struct BitStart;

	template <>
	struct BitStart<>
	{
		size_t raw;
		constexpr BitStart(size_t v) noexcept : raw(v) {}
		constexpr operator size_t() const { return raw; }
	};

	template <size_t START>
	struct BitStart<START>
	{
		constexpr operator size_t() const { return START; }
	};

	template <typename FLD, size_t START, std::unsigned_integral UNDER>
	constexpr FLD get_fld(UNDER under, BitStart<START> = {})
	{
		using field_value_type = typename FldTraits<FLD>::raw;
		using map_value_type = std::conditional_t<std::is_signed_v<field_value_type>, std::make_signed_t<UNDER>, UNDER>;
		constexpr size_t ERemainBits = Math::BitCntOf<UNDER> - FldTraits<FLD>::EBits;
		return FldTraits<FLD>::raw_to_fld(map_value_type(under << ERemainBits - START) >> ERemainBits);
	}

	template <typename FLD, size_t START, std::unsigned_integral UNDER>
	constexpr UNDER with_fld(FLD v, UNDER under, BitStart<START> = {})
	{
		constexpr size_t ERemainBits = Math::BitCntOf<UNDER> - FldTraits<FLD>::EBits;
		constexpr UNDER mask = UNDER(~(UNDER(EAllOne) >> ERemainBits << START));
		return (under & mask) | (UNDER(FldTraits<FLD>::fld_to_raw(v)) << START);
	}

	template <typename FLD, std::unsigned_integral UNDER>
	constexpr FLD get_fld(UNDER under, BitStart<> start)
	{
		using field_value_type = typename FldTraits<FLD>::raw;
		using map_value_type = std::conditional_t<std::is_signed_v<field_value_type>, std::make_signed_t<UNDER>, UNDER>;
		constexpr size_t ERemainBits = Math::BitCntOf<UNDER> - FldTraits<FLD>::EBits;
		return FldTraits<FLD>::raw_to_fld(map_value_type(under << ERemainBits - start) >> ERemainBits);
	}

	template <typename FLD, std::unsigned_integral UNDER>
	constexpr UNDER with_fld(FLD v, UNDER under, BitStart<> start)
	{
		constexpr size_t ERemainBits = Math::BitCntOf<UNDER> - FldTraits<FLD>::EBits;
		UNDER mask = ~(UNDER(EAllOne) >> ERemainBits << start);
		return (under & mask) | (UNDER(FldTraits<FLD>::fld_to_raw(v)) << start);
	}
} // namespace pw_pbl::_bitfield

namespace pw_pbl::inline bitfield
{
	/// 位域引用
	/// @tparam UNDER 位域绑定的实际底层对象类型
	/// @tparam FLD 位域类型，隐含决定了位域宽度
	/// @tparam START 位域在底层的起始位，不指定时表示可变，引用可绑定不同起始位的位域
	template <typename UNDER, typename FLD, size_t... START>
	struct BitFieldRef;

	template <>
	struct BitFieldStruct<>
	{
	protected:
		constexpr BitFieldStruct() = default;
	};

	template <ccMaybeUnalignUnsignedIntegral UNDER>
	struct BitFieldStruct<UNDER> : BitFieldStruct<>
	{
		using raw_type = UNDER;
		using underlying_type = MaybeUnalignTraits<raw_type>::raw_type;

		raw_type raw;

		constexpr BitFieldStruct() = default;
		constexpr BitFieldStruct(underlying_type v) noexcept : raw(v) {}

		template <typename SELF> requires(sizeof(SELF) == sizeof(BitFieldStruct))
		constexpr bool operator==(this const SELF &self, const SELF &oth)
		{
			return MaybeUnalignTraits<>::read(static_cast<const BitFieldStruct>(self).raw) == MaybeUnalignTraits<>::read(static_cast<const BitFieldStruct>(oth).raw);
		}
	};

	/// 位域引用静态抽象基类
	/// @tparam FLD 位域类型，隐含决定了位域宽度
	template <typename FLD>
	struct BitFieldRef<void, FLD>
	{
	protected:
		constexpr BitFieldRef() = default;
	};

	/// 位域引用
	/// @tparam UNDER 位域绑定的实际位结构
	/// @tparam FLD 位域类型，隐含决定了位域宽度
	/// @tparam START 位域的起始位
	template <typename UNDER, _bitfield::ccAsBitField FLD, size_t... START> requires(sizeof...(START) <= 1)
	struct BitFieldRef<UNDER, FLD, START...> : BitFieldRef<void, FLD>
	{
		using under_traits = _bitfield::UnderTraits<UNDER>;
		using raw_ref_type = typename under_traits::raw_ref_type;
		using struct_type = typename under_traits::struct_type;
		/// 去掉所有结构和cvr修饰后最底层的整型类型
		using underlying_type = typename under_traits::underlying_type;
		using field_type = FLD;
		static constexpr bool EIsDynamicStart = sizeof...(START) == 0;
		static constexpr size_t EConstStart = (0 + ... + START);
		static constexpr auto EUnderBits = Math::BitCntOf<underlying_type>;
		static_assert(EConstStart < EUnderBits);
		static constexpr auto EBits = _bitfield::FldTraits<FLD>::EBits;
		static_assert(EConstStart + EBits <= EUnderBits);

		raw_ref_type raw;
		[[no_unique_address]] _bitfield::BitStart<START...> bit_start;

		constexpr BitFieldRef(raw_ref_type under, _bitfield::BitStart<START...> = {}) noexcept requires(!EIsDynamicStart) : raw{std::forward<raw_ref_type>(under)} {}
		constexpr BitFieldRef(raw_ref_type under, size_t st) noexcept requires EIsDynamicStart : raw{std::forward<raw_ref_type>(under)}, bit_start(st) {}

		/// !原始指针可能没有对齐
		constexpr auto *raw_ptr() const { return under_traits::raw_ptr(raw); }

		constexpr FLD read() const { return _bitfield::get_fld<FLD>(under_traits::read(raw), bit_start); }

		constexpr struct_type with(FLD v) const requires(!std::is_void_v<struct_type>) { return {_bitfield::with_fld(std::move(v), under_traits::read(raw), bit_start)}; }

		constexpr raw_ref_type write(FLD v) const requires(!under_traits::EIsConst)
		{
			under_traits::alter(raw, [v = std::move(v), bit_start = this->bit_start](underlying_type &ref) { ref = _bitfield::with_fld(std::move(v), ref, bit_start); });
			return std::forward<raw_ref_type>(raw);
		}

		template <typename F>
		constexpr raw_ref_type write(F f) requires(!under_traits::EIsConst) && std::invocable<F, FLD &>
		{
			under_traits::alter(raw, [f = std::move(f), bit_start = this->bit_start](underlying_type &ref) {
				auto fld = _bitfield::get_fld<FLD>(ref, bit_start);
				std::move(f)(fld);
				ref = _bitfield::with_fld(std::move(fld), ref, bit_start);
			});
			return std::forward<raw_ref_type>(raw);
		}

		explicit(std::same_as<FLD, bool>) constexpr operator FLD() const { return read(); }

		template <typename TGT_FLD> requires(_bitfield::FldTraits<TGT_FLD>::EBits == EBits)
		constexpr BitFieldRef<UNDER, TGT_FLD, START...> cast_to() const { return {raw, bit_start}; }
	};

/// 在位结构内定义位域
/// @param name 名字
/// @param START 起始位
/// @param ... 位域类型（宏逗号问题所以使用...替代名字）
#define PW_PBL_BITFIELD(name, START, ...) \
	constexpr auto name(this auto &&self) -> BitFieldRef<decltype(self), __VA_ARGS__, START> { return {self}; }

	/// 位域数组引用
	/// @tparam UNDER 位域绑定的实际底层对象类型
	/// @tparam FLD 位域类型，隐含决定了位域宽度
	/// @tparam N 数组元素个数
	/// @tparam START 位域在底层的起始位，不指定时表示可变，引用可绑定不同起始位的位域
	template <typename UNDER, typename FLD, size_t N, size_t... START> requires(sizeof...(START) <= 1)
	struct BitFieldRef<UNDER, FLD[N], START...>
	{
		using element_type = BitFieldRef<UNDER, FLD, START...>;
		using raw_ref_type = typename element_type::raw_ref_type;
		static_assert((START + ... + 0) + N * element_type::EBits <= element_type::EUnderBits);

		element_type _first;

		explicit constexpr BitFieldRef(element_type first) : _first(first) {}
		constexpr BitFieldRef(raw_ref_type under, _bitfield::BitStart<START...> = {}) noexcept requires(sizeof...(START) != 0) : _first{std::forward<raw_ref_type>(under)} {}
		constexpr BitFieldRef(raw_ref_type under, size_t st) noexcept requires(sizeof...(START) == 0) : _first{std::forward<raw_ref_type>(under), st} {}

		constexpr auto operator[](size_t i) const { return BitFieldRef<UNDER, FLD>{_first.raw, _first.bit_start + i * element_type::EBits}; }

		template <typename INDEX> requires std::is_enum_v<INDEX>
		constexpr auto operator[](INDEX i) const { return BitFieldRef<UNDER, FLD>{_first.raw, uint8_t(_first.bit_start + std::to_underlying(i) * element_type::EBits)}; }

		template <size_t INDEX> requires(sizeof...(START) == 1 && INDEX < N)
		constexpr auto index() const { return BitFieldRef<UNDER, FLD, (START + ... + 0) + INDEX * element_type::EBits>{_first.raw}; }

		template <auto INDEX> requires(sizeof...(START) == 1 && std::to_underlying(INDEX) < N)
		constexpr auto index() const { return BitFieldRef<UNDER, FLD, (START + ... + 0) + std::to_underlying(INDEX) * element_type::EBits>{_first.raw}; }

		template <typename TGT_FLD> requires(_bitfield::FldTraits<std::remove_extent_t<TGT_FLD>>::EBits == element_type::EBits && std::extent_v<TGT_FLD> <= N)
		constexpr auto cast_to() const { return BitFieldRef<UNDER, TGT_FLD, START...>{_first.template cast_to<std::remove_extent_t<TGT_FLD>>()}; }
	};

} // namespace pw_pbl::inline bitfield

namespace pw_pbl
{
	template <typename UNDER, typename FLD, size_t... START>
	class [[deprecated("undone")]] Option<BitFieldRef<UNDER, FLD, START...>> : public Option<std::type_identity<BitFieldRef<UNDER, FLD, START...>>>
	{
		BitFieldRef<UNDER, FLD, START...> raw;

	public:
		constexpr Option() : raw(nullptr) {}
		constexpr Option(BitFieldRef<UNDER, FLD, START...> data) : raw(data) {}
		constexpr bool is_none() const { return raw.raw == nullptr; }
		constexpr const BitFieldRef<UNDER, FLD, START...> &get_unchecked() const { return raw; }
	};

} // namespace pw_pbl

#endif // PW_PBL_BITFIELD_HPP
